from fastapi import APIRouter, Depends, Header, Body
from fastapi.responses import StreamingResponse

import typing as T
from ...pojo import Result, ChatRequestDTO, ChatPlusRequestDTO, GeneratorImageDTO
from ..service.impl.ChatServiceImpl import ChatService

router = APIRouter(prefix="/llm", tags=["大模型相关功能接口"])


@router.post("/chat")
async def chat(
    chatService: ChatService,
    chatRequest: T.Annotated[ChatRequestDTO, Body(description="聊天的基本参数")],
):
    """
    # 大模型聊天【普通】

    - 普通的对话聊天，如果要没有传递会话id, 和用户id在请求头中就不会带有历史聊天记录。

    - 支持流式响应，将请求体中的stream参数设置为true即可

    ## 请求示例

    ```json
    {
        "question": "你好",
        "model": "qwen-plus", // 模型名称【可选】
        "sessionId": 1, // 会话id，如果在请求头中没填就不会给模型历史对话记录
        "stream": false, // 是否流式推理【可选】
        "temperature": 0.8, // 温度，用于控制生成结果的随机性，值越大表示结果更随机，值越小表示结果越稳定，建议选择在0~2之间【可选】
        "top_p": 0.4, // 核采样，控制生成文本随机性的，值越大表示结果越随机，值越小表示结果越稳定，0~1之间【可选】
        "max_completion_tokens": 8672, // 最大生成token数【可选】
        "frequency_penalty": 0.5, // 频率惩罚，用于控制模型生成的结果中，重复出现的词的频率，值越大表示结果越随机，值越小表示结果越稳定，-2~2之间【可选】
        "presence_penalty": 0.5, // 存在惩罚，用于控制模型生成的结果中，新出现的词的频率，值越大表示结果越随机，值越小表示结果越稳定，-2~2之间【可选】
        "modalities": ["text"] // 模型输出的模态目前只有['text']，默认为['text']【建议别填】
    }
    ```
    ## 返回示例【非流式】

    ```json
    {
        "code": 1,
        "msg": "success",
        "data":
            "你好！很高兴为你提供教学备课支持。请告诉我你需要备课的学科、年级、主题以及具体需求，我会根据你的要求制定详细的 教案和教学建议。
            例如：
            - **学科**：
                语文、数学、英语、物理、化学等- **年级**：小学、初中、高中或大学
            - **主题**：
                具体的课程内容或知识点
            - **需求**：
                教案、课件、互动设计、作业设计等期待你的详细信息！"
    }
    ```

    """
    data = await chatService.chat(chatRequest)
    return (
        Result.success(data)
        if not chatRequest.stream
        else StreamingResponse(data, media_type="text/event-stream")
    )


@router.post("/chatPlus")
async def chatplus(
    chatService: ChatService,
    chatPlusRequest: T.Annotated[
        ChatPlusRequestDTO, Body(description="知识检索的参数")
    ],
    knowledgeBaseId: T.Annotated[int, Header(description="知识库的id[必填]")],
):
    """
    # 大模型聊天【高级】

    带有知识检索的对话聊天
    - 必须要要传递知识库id到请求头中
    - 其他参数同普通聊天

    ## 请求示例

    ```json
    {
        "question": "你好",
        "model": "qwen-plus", // 模型名称，默认为qwen-plus
        "sessionId": 1, // 会话id，如果在请求头中没填就不会给模型历史对话记录
        "knowledgeBaseId": 1, //请求头中必填
        "stream": false // 是否流式响应，默认为false
    }
    ```

    ## 返回示例【非流式】

    略.
    """
    chatPlusRequest.knowledgeBaseId = knowledgeBaseId
    data = await chatService.chatPlus(chatPlusRequest)
    return (
        Result.success(data)
        if not chatPlusRequest.stream
        else StreamingResponse(data, media_type="text/event-stream")
    )


@router.post("/image", response_model=Result)
async def generator_image(
    server: ChatService,
    generatorImageDTO: T.Annotated[GeneratorImageDTO, Body(description="图片生成参数")],
):
    """
    # 图片生成

    更具描述生成图片, 参数介绍

    ---

    ### **GeneratorImageDTO 模型字段说明**

    #### **1. Input 子模型**

    | 字段名          | 类型          | 必填 | 默认值 | 描述                                                   | 约束             |
    | --------------- | ------------- | ---- | ------ | ------------------------------------------------------ | ---------------- |
    | prompt          | str           | 是   | 无     | 正向提示词，描述生成图像中期望包含的元素和视觉特点。   | `max_length=800` |
    | negative_prompt | Optional[str] | 否   | None   | 反向提示词，描述生成图像中不希望包含的元素和视觉特点。 | `max_length=500` |

    ---

    #### **2. Parameters 子模型**

    | 字段名        | 类型           | 必填 | 默认值       | 描述                                                                      | 约束                                                        |
    | ------------- | -------------- | ---- | ------------ | ------------------------------------------------------------------------- | ----------------------------------------------------------- |
    | size          | Optional[str]  | 否   | "1024\*1024" | 图像宽高边长的像素范围为 `[768, 1440]`。格式如 `"1024*1024"`。            | 需符合 `width*height` 格式，且数值在 `[768, 1440]` 范围内。 |
    | n             | Optional[int]  | 否   | 1            | 生成图像的数量，默认为 `1`，最大值为 `4`。                                | `ge=1`, `le=4`（即 `1 <= n <= 4`）                          |
    | seed          | Optional[int]  | 否   | None         | 随机数种子，控制生成内容的随机性。取值范围 `[0, 2147483647]`。            | `ge=0`, `le=2147483647`                                     |
    | prompt_extend | Optional[bool] | 否   | False        | 是否开启 prompt 智能改写（仅对正向提示词生效）。开启后会增加 3-4 秒耗时。 | 无约束，但开启需注意性能影响。                              |
    | watermark     | Optional[bool] | 否   | False        | 是否添加水印。                                                            | 无约束                                                      |

    ---

    #### **3. GeneratorImageDTO 主模型**

    | 字段名     | 类型                         | 必填 | 默认值              | 描述                                                     | 约束                                         |
    | ---------- | ---------------------------- | ---- | ------------------- | -------------------------------------------------------- | -------------------------------------------- |
    | input      | Input                        | 是   | 无                  | 包含对图片的描述（prompt 和 negative_prompt）。          | `prompt` 必填，且需符合 `Input` 子模型约束。 |
    | model      | Literal["wanx2.1-t2i-turbo"] | 是   | "wanx2.1-t2i-turbo" | 模型名称，目前仅支持 `wanx2.1-t2i-turbo`。               | 固定值，不可修改。                           |
    | parameters | Optional[Parameters]         | 否   | `Parameters()`      | 图像生成参数（如尺寸、数量、种子等），默认使用默认参数。 | 需符合 `Parameters` 子模型约束。             |

    **返回的的data里面是url列表**

    ## 请求示例
    ```json
    {
        "input": {
            "prompt": "蓝天，白云, 平静的湖面"
        },
        "model": "wanx2.1-t2i-turbo",
        "parameters": {
                "size": "1024*1024",
                "n": 1,
                "prompt_extend": false,
                "watermark": false
            }
    }
    ```

    ## 返回示例

    ```json
    {
        "code": 1,
        "msg": "success",
        "data": [
            "https://.../xxx.png"
        ]
    }
    ```

    """
    tag, result = await server.generator_image(generatorImageDTO)
    return Result.success(result) if tag == True else Result.error(tag, result)
